{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ANALOGY 实践\n",
    "\n",
     "by 张文 [wenzhang2015@zju.edu.cn](wenzhang2015@zju.edu.cn) \n",
    “\n”,
    "在这个演示中，我们使用ANALOGY([论文链接](https://arxiv.org/pdf/1705.02426.pdf))对示例中文知识图谱进行链接预测，从而达到补全知识图谱的目的。\n",
    "\n",
    "希望在这个demo中帮助大家了解知识图谱表示学习的作用原理和机制。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据集\n",
    "这个示例中，我们使用的是表示学习模型做知识图谱链接预测常用的benchmark数据集 FB15k-237：\n",
    "\n",
    "| #Ent | #Rel | # Train | #Test | #Valid |\n",
    "| --- | --- | --- | --- | --- |\n",
    "| 14,541| 237| 272,115| 17,535| 20,466 |\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ANALOGY 原理回顾\n",
    "ANALOGY提出是为了提升表示学习方法的类比推理能力，采用了DistMult的线性变换假设，将实体表示为实数空间的向量，将关系表示为实数空间的矩阵。每个三原则的得分函数设计如下：\n",
    "$$ \\phi (h,r,t) = \\mathbf{h}^\\top \\mathbf{M}_r \\mathbf{t}$$\n",
    "损失函数定义如下：\n",
    "$$ L = \\sum_{(h,r,t, label)\\in D} -log\\sigma(lable*\\phi (h,r,t))) $$\n",
    "其中D是包含正样本和负样本的数据集合，正样本的lable为1，负样本的lable为-1，$\\sigma$是sigmoid激活函数。\n",
    "\n",
    "类比推理的基本结构如下图：\n",
    "<img src=\"./pic/ANALOGY.png\" style=\"zoom:50%;\" />\n",
    "通过对类比推理的基本结构的分析，ANALOGY中在优化目标损失函数中增加了两个约束如下：\n",
    "$$ \\mathbf{M}_r \\mathbf{M}_r^\\top = \\mathbf{M}_r^\\top \\mathbf{M}_r \\qquad \\forall r \\in \\mathcal{R} $$\n",
    "$$ \\mathbf{M}_r \\mathbf{M}_{r^\\prime} = \\mathbf{M}_{r^\\prime} \\mathbf{M}_r \\qquad \\forall r \\in \\mathcal{R} $$\n",
    "由于存在公理：任意一个正规矩阵$\\mathbf{A}$,一定存在一个正交矩阵$\\mathbf{Q}$和一个分块对角矩阵$\\mathbf{B}$是的 $\\mathbf{A} = \\mathbf{Q}\\mathbf{B}\\mathbf{Q}^\\top$, 并且$\\mathbf{B}$的每个分块里的形式如(1)一个实数，或(2)一个形如 ：\n",
    "\\begin{vmatrix}\n",
    "        x & -y \\\\\n",
    "        y & x \\\\\n",
    "        \\end{vmatrix} \n",
    "的二维矩阵，其中$x,y$均为实数。\n",
    "\n",
    "如果$\\mathbf{M}_r = \\mathbf{Q}\\mathbf{B}\\mathbf{Q}^\\top$, 那么$\\phi (h,r,t) = \\mathbf{h}^\\top \\mathbf{Q}\\mathbf{B}\\mathbf{Q}^\\top \\mathbf{t}$, 我们可以重新令$\\mathbf{h} <- \\mathbf{Q}^\\top\\mathbf{h}$，$\\mathbf{t} <- \\mathbf{Q}^\\top\\mathbf{t}$, 并且$\\mathbf{M}_r = \\mathbf{B}$. 所以ANALOGY最终将知识图谱中的实体表示为了符合约束的分块对角矩阵。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 代码实践\n",
    "在本时间中，我们假设将关系表示为一个$d \\times d$的分块对角矩阵，其左上角的$\\frac{d}{2} \\times \\frac{d}{2}$的分块里为一个对角矩阵，其右下角$\\frac{d}{2} \\times \\frac{d}{2}$的分块为一个分块对角矩阵，且每个小分块满足形式：\n",
    "\\begin{vmatrix}\n",
    "        x & -y \\\\\n",
    "        y & x \\\\\n",
    "        \\end{vmatrix} \n",
    "所以$d$必须是4的倍数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/zhangwen/anaconda3/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
      "  from ._conv import register_converters as _register_converters\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf \n",
    "import time \n",
    "import argparse\n",
    "import random\n",
    "import numpy as np \n",
    "import os.path\n",
    "import math\n",
    "import timeit\n",
    "from multiprocessing import JoinableQueue, Queue, Process\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DistMult:\n",
    "    @property\n",
    "    def variables(self):\n",
    "        return self.__variables\n",
    "\n",
    "    @property\n",
    "    def num_triple_train(self):\n",
    "        return self.__num_triple_train\n",
    "\n",
    "    @property \n",
    "    def num_triple_test(self):\n",
    "        return self.__num_triple_test\n",
    "\n",
    "    @property\n",
    "    def testing_data(self):\n",
    "        return self.__triple_test\n",
    "\n",
    "    @property \n",
    "    def num_entity(self):\n",
    "        return self.__num_entity\n",
    "\n",
    "    @property\n",
    "    def embedding_entity(self):\n",
    "        return self.__embedding_entity\n",
    "\n",
    "\n",
    "    @property\n",
    "    def embedding_relation(self):\n",
    "        return self.__embedding_relation\n",
    "\n",
    "    @property\n",
    "    def hr_t(self):\n",
    "        return self.__hr_t\n",
    "\n",
    "    @property \n",
    "    def tr_h(self):\n",
    "        return self.__tr_h\n",
    "    \n",
    "    @property\n",
    "    def entity2id(self):\n",
    "        return self.__entity2id\n",
    "    \n",
    "    @property\n",
    "    def relation2id(self):\n",
    "        return self.__relation2id\n",
    "\n",
    "    @property\n",
    "    def id2entity(self):\n",
    "        return self.__id2entity\n",
    "    \n",
    "    @property\n",
    "    def id2relation(self):\n",
    "        return self.__id2relation\n",
    "\n",
    "    def training_data_batch(self, batch_size = 512):\n",
    "        n_triple = len(self.__triple_train)\n",
    "        rand_idx = np.random.permutation(n_triple)\n",
    "        start = 0\n",
    "        while start < n_triple:\n",
    "            start_t = timeit.default_timer()\n",
    "            end = min(start+batch_size, n_triple)\n",
    "            size = end - start \n",
    "            train_triple_positive = np.asarray([ self.__triple_train[x] for x in  rand_idx[start:end]])\n",
    "\n",
    "            num_neg = 1\n",
    "            train_negative1 = np.repeat(train_triple_positive, num_neg, axis=0)\n",
    "            train_negative2 = np.repeat(train_triple_positive, num_neg, axis=0)\n",
    "            train_negative1[:, 0] = np.random.randint(self.__num_entity, size=num_neg*size)\n",
    "            train_negative2[:, 2] = np.random.randint(self.__num_entity, size=num_neg*size)\n",
    "            train_triple_negative = np.concatenate((train_negative1, train_negative2), axis=0)\n",
    "\n",
    "            start = end\n",
    "            prepare_t = timeit.default_timer()-start_t\n",
    "\n",
    "            yield train_triple_positive, train_triple_negative, prepare_t\n",
    "\n",
    "\n",
    "    def __init__(self, data_dir, negative_sampling,learning_rate, \n",
    "             batch_size, max_iter, margin, dimension, norm, evaluation_size, regularizer_weight):\n",
    "        # this part for data prepare\n",
    "        self.__data_dir=data_dir\n",
    "        self.__negative_sampling=negative_sampling\n",
    "        self.__regularizer_weight = regularizer_weight\n",
    "        self.__norm = norm\n",
    "\n",
    "        self.__entity2id={}\n",
    "        self.__id2entity={}\n",
    "        self.__relation2id={}\n",
    "        self.__id2relation={}\n",
    "\n",
    "        self.__triple_train=[] #[(head_id, relation_id, tail_id),...]\n",
    "        self.__triple_test=[]\n",
    "        self.__triple_valid=[]\n",
    "        self.__triple = []\n",
    "\n",
    "        self.__num_entity=0\n",
    "        self.__num_relation=0\n",
    "        self.__num_triple_train=0\n",
    "        self.__num_triple_test=0\n",
    "        self.__num_triple_valid=0\n",
    "\n",
    "        # load all the file: entity2id.txt, relation2id.txt, train.txt, test.txt, valid.txt\n",
    "        self.load_data()\n",
    "        print('finish preparing data. ')\n",
    "\n",
    "\n",
    "        # this part for the model:\n",
    "        self.__learning_rate = learning_rate\n",
    "        self.__batch_size = batch_size\n",
    "        self.__max_iter = max_iter\n",
    "        self.__margin = margin\n",
    "        self.__dimension = dimension\n",
    "        self.dim = dimension\n",
    "        self.__variables= []\n",
    "        #self.__norm = norm\n",
    "        self.__evaluation_size = evaluation_size\n",
    "        bound = 6 / math.sqrt(self.__dimension)\n",
    "        with tf.device('/cpu'):\n",
    "            self.__embedding_entity = tf.get_variable('embedding_entity', [self.__num_entity, self.__dimension],\n",
    "                                                       initializer=tf.random_uniform_initializer(minval=-bound, maxval=bound, seed = 123), dtype=tf.float32)\n",
    "            self.__embedding_relation = tf.get_variable('embedding_relation', [self.__num_relation, self.__dimension],\n",
    "                                                         initializer=tf.random_uniform_initializer(minval=-bound, maxval=bound, seed =124), dtype=tf.float32)\n",
    "            self.__variables.append(self.__embedding_entity)\n",
    "            self.__variables.append(self.__embedding_relation)\n",
    "            print('finishing initializing')\n",
    "\n",
    "\n",
    "    def load_data(self):\n",
    "        print('loading entity2id.txt ...')\n",
    "        with open(os.path.join(self.__data_dir, 'entity2id.txt'), encoding='utf-8') as f:\n",
    "            self.__entity2id = {line.strip().split('\\t')[0]: int(line.strip().split('\\t')[1]) for line in f.readlines()}\n",
    "            self.__id2entity = {value:key for key,value in self.__entity2id.items()}\n",
    "\n",
    "        print('loading reltion2id.txt ...')     \n",
    "        with open(os.path.join(self.__data_dir,'relation2id.txt'), encoding='utf-8') as f:\n",
    "            self.__relation2id = {line.strip().split('\\t')[0]: int(line.strip().split('\\t')[1]) for line in f.readlines()}\n",
    "            self.__id2relation = {value:key for key, value in self.__relation2id.items()}\n",
    "\n",
    "        def load_triple(self, triplefile):\n",
    "            triple_list = [] #[(head_id, relation_id, tail_id),...]\n",
    "            with open(os.path.join(self.__data_dir, triplefile), encoding='utf-8') as f:\n",
    "                for line in f.readlines():\n",
    "                    line_list = line.strip().split('\\t')\n",
    "                    assert len(line_list) == 3\n",
    "                    headid = self.__entity2id[line_list[0]]\n",
    "                    relationid = self.__relation2id[line_list[1]]\n",
    "                    tailid = self.__entity2id[line_list[2]]\n",
    "                    triple_list.append((headid, relationid, tailid))\n",
    "                    self.__hr_t[(headid, relationid)].add(tailid)\n",
    "                    self.__tr_h[(tailid, relationid)].add(headid)\n",
    "            return triple_list\n",
    "\n",
    "        self.__hr_t = defaultdict(set)\n",
    "        self.__tr_h = defaultdict(set)\n",
    "        self.__triple_train = load_triple(self, 'train.txt')\n",
    "        self.__triple_test = load_triple(self, 'test.txt')\n",
    "        self.__triple_valid = load_triple(self, 'valid.txt')\n",
    "        self.__triple = np.concatenate([self.__triple_train, self.__triple_test, self.__triple_valid], axis = 0 )\n",
    "\n",
    "        self.__num_relation = len(self.__relation2id)\n",
    "        self.__num_entity = len(self.__entity2id)\n",
    "        self.__num_triple_train = len(self.__triple_train)\n",
    "        self.__num_triple_test = len(self.__triple_test)\n",
    "        self.__num_triple_valid = len(self.__triple_valid)\n",
    "\n",
    "        print('entity number: ' + str(self.__num_entity))\n",
    "        print('relation number: ' + str(self.__num_relation))\n",
    "        print('training triple number: ' + str(self.__num_triple_train))\n",
    "        print('testing triple number: ' + str(self.__num_triple_test))\n",
    "        print('valid triple number: ' + str(self.__num_triple_valid))\n",
    "\n",
    "\n",
    "        if self.__negative_sampling == 'bern':\n",
    "            self.__relation_property_head = {x:[] for x in range(self.__num_relation)} #{relation_id:[headid1, headid2,...]}\n",
    "            self.__relation_property_tail = {x:[] for x in range(self.__num_relation)} #{relation_id:[tailid1, tailid2,...]}\n",
    "            self.__relation_property = {x:[] for x in range(self.__num_relation)} \n",
    "            for t in self.__triple_train:\n",
    "                #print(t)\n",
    "                self.__relation_property_head[t[1]].append(t[0])\n",
    "                self.__relation_property_tail[t[1]].append(t[2])\n",
    "            #print(self.__relation_property_head[0])\n",
    "            #print(self.__relation_property_tail[0])\n",
    "            for x in self.__relation_property_head.keys():\n",
    "                t = len(set(self.__relation_property_tail[x]))\n",
    "                h = len(set(self.__relation_property_head[x]))\n",
    "                self.__relation_property[x] = float(t)/(h+t+0.000000001)\n",
    "            #self.__relation_property = {x:(len(set(self.__relation_property_tail[x])))/(len(set(self.__relation_property_head[x]))+ len(set(self.__relation_property_tail[x]))) \\\n",
    "            #\t\t\t\t\t\t\t for x in self.__relation_property_head.keys()} # {relation_id: p, ...} 0< num <1, and for relation replace head entity with the property p\n",
    "        else: \n",
    "            print(\"unif set don't need to calculate hpt and tph\")\n",
    "\n",
    "    def split_embedding(self, embedding):\n",
    "        assert self.dim % 4 == 0\n",
    "        num_scalar = self.dim // 2\n",
    "        num_block = self.dim // 4\n",
    "        embedding_scalar = embedding[:, 0:num_scalar]\n",
    "        embedding_x = embedding[:, num_scalar:-num_block]\n",
    "        embedding_y = embedding[:, -num_block:]\n",
    "        return embedding_scalar, embedding_x, embedding_y\n",
    "    \n",
    "    def score(self, h,r,t):\n",
    "        h_scalar, h_x ,h_y = self.split_embedding(h)\n",
    "        r_scalar, r_x, r_y = self.split_embedding(r)\n",
    "        t_scalar, t_x, t_y = self.split_embedding(t)\n",
    "        score_scalar = tf.reduce_sum(h_scalar * r_scalar * t_scalar, axis=1)\n",
    "        print(h_scalar, h_x ,h_y)\n",
    "        score_block = tf.reduce_sum(h_x * r_x * t_x\n",
    "                                    + h_x * r_y * t_y\n",
    "                                    + h_y * r_x * t_y\n",
    "                                    - h_y * r_y * t_x, axis=1)\n",
    "        score = score_scalar + score_block\n",
    "        return score\n",
    "    \n",
    "\n",
    "    def train(self, inputs):\n",
    "        embedding_relation = self.__embedding_relation\n",
    "        embedding_entity = self.__embedding_entity\n",
    "\n",
    "        triple_positive, triple_negative = inputs # triple_positive:(head_id,relation_id,tail_id)\n",
    "\n",
    "        norm_entity = embedding_entity\n",
    "        norm_relation = embedding_relation\n",
    "        norm_entity_l2sum = tf.sqrt(tf.reduce_sum(norm_entity**2, axis = 1))\n",
    "\n",
    "        embedding_positive_head = tf.nn.embedding_lookup(norm_entity, triple_positive[:, 0])\n",
    "        embedding_positive_tail = tf.nn.embedding_lookup(norm_entity, triple_positive[:, 2])\n",
    "        embedding_positive_relation = tf.nn.embedding_lookup(norm_relation, triple_positive[:, 1])\n",
    "\n",
    "        embedding_negative_head = tf.nn.embedding_lookup(norm_entity, triple_negative[:, 0])\n",
    "        embedding_negative_tail = tf.nn.embedding_lookup(norm_entity, triple_negative[:, 2])\n",
    "        embedding_negative_relation = tf.nn.embedding_lookup(norm_relation, triple_negative[:, 1])\n",
    "\n",
    "        score_positive = self.score(embedding_positive_head, embedding_positive_relation, embedding_positive_tail)\n",
    "        score_negative = self.score(embedding_negative_head, embedding_negative_relation, embedding_negative_tail)\n",
    "        \n",
    "        score_pos = -tf.log(tf.clip_by_value(tf.sigmoid(score_positive), 1e-32, 1.0))\n",
    "        score_neg = -tf.log(tf.clip_by_value(tf.sigmoid(-score_negative), 1e-32, 1.0))\n",
    "        score = tf.concat((score_pos, score_neg), axis=0)\n",
    "        \n",
    "        loss_triple = tf.reduce_mean(score)\n",
    "        \n",
    "        self.__loss_regularizer = loss_regularizer = tf.reduce_sum(tf.abs(self.__embedding_relation)) + tf.reduce_sum(tf.abs(self.__embedding_entity))\n",
    "        return loss_triple + loss_regularizer*self.__regularizer_weight,  norm_entity_l2sum\n",
    "\n",
    "    def test(self, inputs):\n",
    "        embedding_relation = self.__embedding_relation\n",
    "        embedding_entity = self.__embedding_entity\n",
    "\n",
    "        triple_test = inputs # (headid, tailid, tailid)\n",
    "        head_vec = tf.nn.embedding_lookup(embedding_entity, triple_test[0])\n",
    "        rel_vec = tf.nn.embedding_lookup(embedding_relation, triple_test[1])\n",
    "        tail_vec = tf.nn.embedding_lookup(embedding_entity, triple_test[2])\n",
    "        \n",
    "        head_vec = tf.reshape(head_vec, [-1, self.dim])\n",
    "        rel_vec = tf.reshape(rel_vec, [-1, self.dim])\n",
    "        tail_vec = tf.reshape(tail_vec, [-1, self.dim])\n",
    "\n",
    "        \n",
    "        score_head = self.score(embedding_entity,  rel_vec, tail_vec)\n",
    "        score_tail = self.score(head_vec, rel_vec,  embedding_entity)\n",
    "        \n",
    "        \n",
    "        # _, id_replace_head = tf.nn.top_k(tf.reduce_sum(embedding_entity * rel_vec * tail_vec, axis=1), k=self.__num_entity)\n",
    "        # _, id_replace_tail = tf.nn.top_k(tf.reduce_sum(head_vec * rel_vec * embedding_entity, axis=1), k=self.__num_entity)\n",
    "        _, id_replace_head = tf.nn.top_k(score_head, k=self.__num_entity)\n",
    "        _, id_replace_tail = tf.nn.top_k(score_tail, k=self.__num_entity)\n",
    "        \n",
    "            \n",
    "        return id_replace_head, id_replace_tail"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_operation(model, learning_rate=0.01, margin=1.0, optimizer_str = 'gradient'):\n",
    "    with tf.device('/cpu'):\n",
    "        train_triple_positive_input = tf.placeholder(tf.int32, [None, 3])\n",
    "        train_triple_negative_input = tf.placeholder(tf.int32, [None, 3])\n",
    "\n",
    "        loss, norm_entity = model.train([train_triple_positive_input, train_triple_negative_input])\n",
    "        if optimizer_str == 'gradient':\n",
    "            optimizer = tf.train.GradientDescentOptimizer(learning_rate = learning_rate)\n",
    "        elif optimizer_str == 'rms':\n",
    "            optimizer = tf.train.RMSPropOptimizer(learning_rate = learning_rate)\n",
    "        elif optimizer_str == 'adam':\n",
    "            optimizer = tf.train.AdamOptimizer(learning_rate = learning_rate)\n",
    "        else:\n",
    "            raise NotImplementedError(\"Dose not support %s optimizer\" %optimizer_str)\n",
    "\n",
    "        grads = optimizer.compute_gradients(loss, model.variables)\n",
    "        op_train = optimizer.apply_gradients(grads)\n",
    "\n",
    "        return train_triple_positive_input, train_triple_negative_input, loss, op_train, norm_entity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_operation(model):\n",
    "    with tf.device('/cpu'):\n",
    "        test_triple = tf.placeholder(tf.int32, [3])\n",
    "        head_rank, tail_rank = model.test(test_triple)\n",
    "        return test_triple, head_rank, tail_rank"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Args:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 测试一个样本函数\n",
    "def test_one_sample(model, trp, session):\n",
    "    t = trp\n",
    "    id_replace_head , id_replace_tail  = session.run([head_rank, tail_rank], {test_triple:t})\n",
    "    hr_t = model.hr_t\n",
    "    tr_h = model.tr_h\n",
    "    \n",
    "    hrank = 0\n",
    "    fhrank = 0\n",
    "    predicted_head_tmp = []\n",
    "    for i in range(len(id_replace_head)):\n",
    "        val = id_replace_head[i]\n",
    "        predicted_head_tmp.append(val)\n",
    "        if val == t[0]:\n",
    "            break\n",
    "        else: \n",
    "            hrank += 1\n",
    "            fhrank += 1 \n",
    "            if val in tr_h[(t[2],t[1])]:\n",
    "                fhrank -= 1\n",
    "                _ = predicted_head_tmp.pop()\n",
    "    predicted_head_tmp = [id_replace_head[i] for i in range(len(id_replace_head))]\n",
    "    \n",
    "\n",
    "    trank = 0\n",
    "    ftrank = 0\n",
    "    predicted_tail_tmp = []\n",
    "    for i in range(len(id_replace_tail)):\n",
    "        val = id_replace_tail[i]\n",
    "        predicted_tail_tmp.append(val)\n",
    "        if val == t[2]:\n",
    "            break\n",
    "        else:\n",
    "            trank += 1\n",
    "            ftrank += 1\n",
    "            if val in hr_t[(t[0], t[1])]:\n",
    "                ftrank -= 1\n",
    "                _ = predicted_tail_tmp.pop()\n",
    "    predicted_tail_tmp = [id_replace_tail[i] for i in range(len(id_replace_tail))]\n",
    "\n",
    "    \n",
    "    return hrank+1, fhrank+1, trank+1, ftrank+1, \\\n",
    "            predicted_head_tmp, predicted_tail_tmp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hit(rank_head, rank_tail, k):\n",
    "    n_test = len(rank_head)\n",
    "    assert len(rank_head) == len(rank_tail)\n",
    "    hit_head = np.sum(np.asarray(np.asarray(rank_head)<=k , dtype=np.float32))/n_test\n",
    "    hit_tail = np.sum(np.asarray(np.asarray(rank_tail)<=k , dtype=np.float32))/n_test\n",
    "    hit = (hit_head + hit_tail)/2.0\n",
    "    return hit_head, hit_tail, hit\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(n_test):\n",
    "    predicted_tail = []\n",
    "    predicted_head = []\n",
    "    np.random.shuffle(testing_data)\n",
    "    for i in range(n_test):\n",
    "        print('[%.2f sec] --- testing[%d/%d]' %(timeit.default_timer()-start, i+1, n_test), end='\\r')\n",
    "        t = testing_data[i]\n",
    "        hrank, fhrank, trank, ftrank,predicted_head_tmp, predicted_tail_tmp = test_one_sample(model, t, session)\n",
    "        rank_head.append(hrank)\n",
    "        rank_tail.append(trank)\n",
    "        filter_rank_head.append(fhrank)\n",
    "        filter_rank_tail.append(ftrank)\n",
    "        \n",
    "\n",
    "        predicted_tail.append(predicted_tail_tmp)\n",
    "        predicted_head.append(predicted_head_tmp)\n",
    "    mean_rank_head = np.sum(rank_head, dtype=np.float32)/n_test\n",
    "    mean_rank_tail = np.sum(rank_tail, dtype=np.float32)/n_test\n",
    "    filter_mean_rank_head = np.sum(filter_rank_head, dtype=np.float32)/n_test\n",
    "    filter_mean_rank_tail = np.sum(filter_rank_tail, dtype=np.float32)/n_test\n",
    "\n",
    "    mean_reciprocal_rank_head = np.sum(1.0/np.asarray(rank_head, dtype=np.float32))/n_test\n",
    "    mean_reciprocal_rank_tail = np.sum(1.0/np.asarray(rank_tail, dtype=np.float32))/n_test\n",
    "    filter_mean_reciprocal_rank_head = np.sum(1.0/np.asarray(filter_rank_head, dtype=np.float32))/n_test\n",
    "    filter_mean_reciprocal_rank_tail = np.sum(1.0/np.asarray(filter_rank_tail, dtype=np.float32))/n_test\n",
    "\n",
    "    hit1_head, hit1_tail, hit1 = hit(rank_head, rank_tail, 1)\n",
    "    filter_hit1_head, filter_hit1_tail, filter_hit1 = hit(filter_rank_head, filter_rank_tail, 1)\n",
    "    hit3_head, hit3_tail, hit3 = hit(rank_head, rank_tail, 3)\n",
    "    filter_hit3_head, filter_hit3_tail, filter_hit3 = hit(filter_rank_head, filter_rank_tail, 3)\n",
    "    hit10_head, hit10_tail, hit10 = hit(rank_head, rank_tail, 10)\n",
    "    filter_hit10_head, filter_hit10_tail, filter_hit10 = hit(filter_rank_head, filter_rank_tail, 10)\n",
    "\n",
    "\n",
    "    print('iter:%d --MR: %.2f  --MRR: %.2f  --hit@1: %.2f   --hit@3: %.2f    --hit@10: %.2f' %(n_iter, (mean_rank_head+ mean_rank_tail)/2, \n",
    "                                                            (mean_reciprocal_rank_head + mean_reciprocal_rank_tail)/2, \n",
    "                                                            hit1, hit3, hit3))\n",
    "    print('iter:%d --FMR: %.2f --FMRR: %.2f --Fhit@1: %.2f  --Fhit@3: %.2f   --Fhit@10: %.2f' %(n_iter, (filter_mean_rank_head+ filter_mean_rank_tail)/2, \n",
    "                                                                    (filter_mean_reciprocal_rank_head + filter_mean_reciprocal_rank_tail)/2,\n",
    "                                                                    filter_hit1, filter_hit3, filter_hit10)) \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<__main__.Args object at 0x13ae3f588>\n",
      "loading entity2id.txt ...\n",
      "loading reltion2id.txt ...\n",
      "entity number: 14541\n",
      "relation number: 237\n",
      "training triple number: 272115\n",
      "testing triple number: 20466\n",
      "valid triple number: 17535\n",
      "finish preparing data. \n",
      "WARNING:tensorflow:From /Users/zhangwen/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Colocations handled automatically by placer.\n",
      "finishing initializing\n",
      "Tensor(\"strided_slice_6:0\", shape=(?, 100), dtype=float32, device=/device:CPU:*) Tensor(\"strided_slice_7:0\", shape=(?, 50), dtype=float32, device=/device:CPU:*) Tensor(\"strided_slice_8:0\", shape=(?, 50), dtype=float32, device=/device:CPU:*)\n",
      "Tensor(\"strided_slice_15:0\", shape=(?, 100), dtype=float32, device=/device:CPU:*) Tensor(\"strided_slice_16:0\", shape=(?, 50), dtype=float32, device=/device:CPU:*) Tensor(\"strided_slice_17:0\", shape=(?, 50), dtype=float32, device=/device:CPU:*)\n",
      "WARNING:tensorflow:From /Users/zhangwen/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/math_ops.py:3066: to_int32 (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.cast instead.\n",
      "Tensor(\"strided_slice_27:0\", shape=(14541, 100), dtype=float32, device=/device:CPU:*) Tensor(\"strided_slice_28:0\", shape=(14541, 50), dtype=float32, device=/device:CPU:*) Tensor(\"strided_slice_29:0\", shape=(14541, 50), dtype=float32, device=/device:CPU:*)\n",
      "Tensor(\"strided_slice_36:0\", shape=(1, 100), dtype=float32, device=/device:CPU:*) Tensor(\"strided_slice_37:0\", shape=(1, 50), dtype=float32, device=/device:CPU:*) Tensor(\"strided_slice_38:0\", shape=(1, 50), dtype=float32, device=/device:CPU:*)\n"
     ]
    }
   ],
   "source": [
    "# 设置参数等\n",
    "args  = Args()\n",
    "args.data_dir = './data/FB15k-237/'\n",
    "args.learning_rate = 0.005\n",
    "args.batch_size = 2048\n",
    "args.optimizer = 'adam'\n",
    "args.dimension = 200\n",
    "args.margin = 3\n",
    "args.norm = 'L2'\n",
    "args.evaluation_size = 500\n",
    "args.save_dir = 'output/'\n",
    "args.negative_sampling = 'bern'\n",
    "args.evaluate_per_iteration = 1\n",
    "args.evaluate_worker = 3\n",
    "args.regularizer_weight = 1e-7\n",
    "args.n_test = 300\n",
    "args.save_per = 100\n",
    "args.n_worker = 5\n",
    "args.max_iter = 20\n",
    "\n",
    "print(args)\n",
    "model = DistMult(negative_sampling=args.negative_sampling, data_dir=args.data_dir,\n",
    "                learning_rate=args.learning_rate, batch_size=args.batch_size,\n",
    "                max_iter=args.max_iter, margin=args.margin, \n",
    "                dimension=args.dimension, norm=args.norm, evaluation_size=args.evaluation_size, \n",
    "                regularizer_weight = args.regularizer_weight)\n",
    "\n",
    "train_triple_positive_input, train_triple_negative_input, loss, op_train, norm_entity = train_operation(model, learning_rate = args.learning_rate, margin = args.margin, optimizer_str = args.optimizer)\n",
    "test_triple, head_rank, tail_rank= test_operation(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/zhangwen/anaconda3/lib/python3.6/site-packages/tensorflow/python/util/tf_should_use.py:193: initialize_all_variables (from tensorflow.python.ops.variables) is deprecated and will be removed after 2017-03-02.\n",
      "Instructions for updating:\n",
      "Use `tf.global_variables_initializer` instead.\n",
      "iter[0] ---loss: 100.62850 ---time: 43.26 ---prepare time : 0.32\n",
      "Model saved at ./save/DistMult_0.ckpt\n",
      "iter:0 --MR: 7174.43  --MRR: 0.00  --hit@1: 0.00   --hit@3: 0.00    --hit@10: 0.00\n",
      "iter:0 --FMR: 7038.70 --FMRR: 0.00 --Fhit@1: 0.00  --Fhit@3: 0.00   --Fhit@10: 0.00\n",
      "iter[1] ---loss: 95.44930 ---time: 39.53 ---prepare time : 0.31\n",
      "iter:1 --MR: 5044.13  --MRR: 0.03  --hit@1: 0.02   --hit@3: 0.03    --hit@10: 0.03\n",
      "iter:1 --FMR: 4908.87 --FMRR: 0.04 --Fhit@1: 0.02  --Fhit@3: 0.03   --Fhit@10: 0.06\n",
      "iter[2] ---loss: 85.90134 ---time: 40.04 ---prepare time : 0.32\n",
      "iter:2 --MR: 2910.73  --MRR: 0.05  --hit@1: 0.03   --hit@3: 0.04    --hit@10: 0.04\n",
      "iter:2 --FMR: 2732.31 --FMRR: 0.08 --Fhit@1: 0.04  --Fhit@3: 0.08   --Fhit@10: 0.16\n",
      "iter[3] ---loss: 71.05738 ---time: 39.58 ---prepare time : 0.31\n",
      "iter:3 --MR: 2164.47  --MRR: 0.05  --hit@1: 0.03   --hit@3: 0.05    --hit@10: 0.05\n",
      "iter:3 --FMR: 2017.15 --FMRR: 0.09 --Fhit@1: 0.05  --Fhit@3: 0.09   --Fhit@10: 0.16\n",
      "iter[4] ---loss: 42.17375 ---time: 39.88 ---prepare time : 0.31\n",
      "iter:4 --MR: 1099.18  --MRR: 0.07  --hit@1: 0.04   --hit@3: 0.06    --hit@10: 0.06\n",
      "iter:4 --FMR: 943.91 --FMRR: 0.13 --Fhit@1: 0.08  --Fhit@3: 0.13   --Fhit@10: 0.21\n",
      "iter[5] ---loss: 27.98954 ---time: 40.16 ---prepare time : 0.31\n",
      "iter:5 --MR: 848.41  --MRR: 0.07  --hit@1: 0.03   --hit@3: 0.07    --hit@10: 0.07\n",
      "iter:5 --FMR: 636.92 --FMRR: 0.12 --Fhit@1: 0.06  --Fhit@3: 0.12   --Fhit@10: 0.22\n",
      "iter[6] ---loss: 22.77698 ---time: 39.46 ---prepare time : 0.31\n",
      "iter:6 --MR: 804.46  --MRR: 0.09  --hit@1: 0.05   --hit@3: 0.09    --hit@10: 0.09\n",
      "iter:6 --FMR: 605.21 --FMRR: 0.15 --Fhit@1: 0.09  --Fhit@3: 0.17   --Fhit@10: 0.26\n",
      "iter[7] ---loss: 20.29978 ---time: 39.21 ---prepare time : 0.30\n",
      "iter:7 --MR: 756.77  --MRR: 0.10  --hit@1: 0.06   --hit@3: 0.10    --hit@10: 0.10\n",
      "iter:7 --FMR: 566.23 --FMRR: 0.16 --Fhit@1: 0.10  --Fhit@3: 0.17   --Fhit@10: 0.28\n",
      "iter[8] ---loss: 18.94841 ---time: 38.58 ---prepare time : 0.30\n",
      "iter:8 --MR: 894.32  --MRR: 0.11  --hit@1: 0.07   --hit@3: 0.11    --hit@10: 0.11\n",
      "iter:8 --FMR: 686.06 --FMRR: 0.18 --Fhit@1: 0.12  --Fhit@3: 0.18   --Fhit@10: 0.30\n",
      "iter[9] ---loss: 17.86277 ---time: 39.17 ---prepare time : 0.32\n",
      "iter:9 --MR: 736.51  --MRR: 0.12  --hit@1: 0.07   --hit@3: 0.12    --hit@10: 0.12\n",
      "iter:9 --FMR: 599.54 --FMRR: 0.20 --Fhit@1: 0.13  --Fhit@3: 0.22   --Fhit@10: 0.34\n",
      "iter[10] ---loss: 17.18015 ---time: 38.71 ---prepare time : 0.30\n",
      "iter:10 --MR: 725.40  --MRR: 0.12  --hit@1: 0.07   --hit@3: 0.12    --hit@10: 0.12\n",
      "iter:10 --FMR: 541.47 --FMRR: 0.20 --Fhit@1: 0.12  --Fhit@3: 0.22   --Fhit@10: 0.36\n",
      "iter[11] ---loss: 16.71703 ---time: 38.59 ---prepare time : 0.30\n",
      "iter:11 --MR: 783.81  --MRR: 0.11  --hit@1: 0.06   --hit@3: 0.11    --hit@10: 0.11\n",
      "iter:11 --FMR: 607.05 --FMRR: 0.19 --Fhit@1: 0.12  --Fhit@3: 0.21   --Fhit@10: 0.34\n",
      "iter[12] ---loss: 16.20879 ---time: 38.64 ---prepare time : 0.30\n",
      "iter:12 --MR: 745.81  --MRR: 0.09  --hit@1: 0.05   --hit@3: 0.09    --hit@10: 0.09\n",
      "iter:12 --FMR: 607.44 --FMRR: 0.17 --Fhit@1: 0.11  --Fhit@3: 0.18   --Fhit@10: 0.31\n",
      "iter[13] ---loss: 15.79837 ---time: 38.57 ---prepare time : 0.30\n",
      "iter:13 --MR: 738.29  --MRR: 0.12  --hit@1: 0.07   --hit@3: 0.12    --hit@10: 0.12\n",
      "iter:13 --FMR: 542.57 --FMRR: 0.20 --Fhit@1: 0.12  --Fhit@3: 0.22   --Fhit@10: 0.34\n",
      "iter[14] ---loss: 15.48467 ---time: 38.63 ---prepare time : 0.30\n",
      "iter:14 --MR: 769.72  --MRR: 0.11  --hit@1: 0.06   --hit@3: 0.12    --hit@10: 0.12\n",
      "iter:14 --FMR: 598.81 --FMRR: 0.18 --Fhit@1: 0.11  --Fhit@3: 0.20   --Fhit@10: 0.35\n",
      "iter[15] ---loss: 15.28255 ---time: 38.67 ---prepare time : 0.30\n",
      "iter:15 --MR: 644.54  --MRR: 0.12  --hit@1: 0.06   --hit@3: 0.12    --hit@10: 0.12\n",
      "iter:15 --FMR: 477.91 --FMRR: 0.21 --Fhit@1: 0.13  --Fhit@3: 0.23   --Fhit@10: 0.37\n",
      "iter[16] ---loss: 14.97615 ---time: 38.82 ---prepare time : 0.30\n",
      "iter:16 --MR: 793.36  --MRR: 0.12  --hit@1: 0.07   --hit@3: 0.11    --hit@10: 0.11\n",
      "iter:16 --FMR: 595.78 --FMRR: 0.20 --Fhit@1: 0.13  --Fhit@3: 0.21   --Fhit@10: 0.34\n",
      "iter[17] ---loss: 14.85544 ---time: 38.60 ---prepare time : 0.30\n",
      "iter:17 --MR: 759.51  --MRR: 0.10  --hit@1: 0.06   --hit@3: 0.10    --hit@10: 0.10\n",
      "iter:17 --FMR: 587.43 --FMRR: 0.18 --Fhit@1: 0.11  --Fhit@3: 0.19   --Fhit@10: 0.33\n",
      "iter[18] ---loss: 14.55947 ---time: 38.72 ---prepare time : 0.31\n",
      "iter:18 --MR: 714.79  --MRR: 0.11  --hit@1: 0.06   --hit@3: 0.10    --hit@10: 0.10\n",
      "iter:18 --FMR: 516.93 --FMRR: 0.19 --Fhit@1: 0.12  --Fhit@3: 0.20   --Fhit@10: 0.35\n",
      "iter[19] ---loss: 14.34992 ---time: 38.52 ---prepare time : 0.30\n",
      "Model saved at ./save/DistMult_19.ckpt\n",
      "iter:19 --MR: 800.42  --MRR: 0.11  --hit@1: 0.06   --hit@3: 0.12    --hit@10: 0.12\n",
      "iter:19 --FMR: 604.28 --FMRR: 0.18 --Fhit@1: 0.10  --Fhit@3: 0.21   --Fhit@10: 0.35\n"
     ]
    }
   ],
   "source": [
    "# 训练模型\n",
    "config = tf.ConfigProto()\n",
    "config.gpu_options.allow_growth = False\n",
    "config.log_device_placement = False\n",
    "config.allow_soft_placement = True\n",
    "config.gpu_options.per_process_gpu_memory_fraction=0.68\n",
    "session = tf.Session(config=config)\n",
    "session.as_default()\n",
    "\n",
    "tf.initialize_all_variables().run(session=session)\n",
    "saver = tf.train.Saver()\n",
    "\n",
    "\n",
    "for n_iter in range(args.max_iter):\n",
    "    accu_loss =0.\n",
    "    batch = 0\n",
    "    num_batch = model.num_triple_train/args.batch_size\n",
    "    start_time = timeit.default_timer()\n",
    "    prepare_time = 0.\n",
    "\n",
    "    for tp, tn , t in  model.training_data_batch(batch_size= args.batch_size):\n",
    "        l, _, norm_e = session.run([loss, op_train, norm_entity], {train_triple_positive_input:tp, train_triple_negative_input: tn})\n",
    "        accu_loss += l\n",
    "        batch += 1\n",
    "        print('[%.2f sec](%d/%d): -- loss: %.5f' %(timeit.default_timer()-start_time, batch, num_batch , l), end='\\r')\n",
    "        prepare_time += t\n",
    "    if n_iter%1 == 0:\n",
    "        print('iter[%d] ---loss: %.5f ---time: %.2f ---prepare time : %.2f' %(n_iter, accu_loss, timeit.default_timer()-start_time, prepare_time))\n",
    "\n",
    "    if n_iter % args.save_per == 0 or n_iter ==0 or n_iter == args.max_iter-1:\n",
    "        save_path = saver.save(session, os.path.join('./save/DistMult_' + str(n_iter) + '.ckpt'))\n",
    "        print('Model saved at %s' % save_path)\n",
    "\n",
    "    if n_iter %args.evaluate_per_iteration == 0 or n_iter ==0 or n_iter == args.max_iter-1:\n",
    "        rank_head = []\n",
    "        rank_tail = []\n",
    "        filter_rank_head = []\n",
    "        filter_rank_tail = []\n",
    "\n",
    "        norm_rank_head = []\n",
    "        norm_rank_tail = []\n",
    "        norm_filter_rank_head = []\n",
    "        norm_filter_rank_tail = []\n",
    "\n",
    "        start = timeit.default_timer()\n",
    "        testing_data = model.testing_data\n",
    "        hr_t = model.hr_t\n",
    "        tr_h = model.tr_h\n",
    "        n_test = args.n_test\n",
    "        test(n_test)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # 测试模型\n",
    "# predicted_tail = []\n",
    "# norm_predicted_tail = []\n",
    "# predicted_head = []\n",
    "# norm_predicted_head = []\n",
    "\n",
    "# rank_head = []\n",
    "# rank_tail = []\n",
    "# filter_rank_head = []\n",
    "# filter_rank_tail = []\n",
    "\n",
    "# norm_rank_head = []\n",
    "# norm_rank_tail = []\n",
    "# norm_filter_rank_head = []\n",
    "# norm_filter_rank_tail = []\n",
    "\n",
    "# start = timeit.default_timer()\n",
    "# testing_data = model.testing_data\n",
    "# # hr_t = model.hr_t\n",
    "# # tr_h = model.tr_h\n",
    "# n_test = args.n_test\n",
    "# if n_iter == args.max_iter-1:\tn_test = model.num_triple_test\n",
    "# predicted_tail = []\n",
    "# norm_predicted_tail = []\n",
    "# predicted_head = []\n",
    "# norm_predicted_head = []\n",
    "# test(n_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们按照设置的参数跑了20个iteration，可以看到ANALOGY在FB15k-237的数据集上已经有了明显的预测效果。\n",
    "本demo中不包括调参的部分，有兴趣的同学可以阅读原文并自行尝试不同的参数组合，并观察对模型训练和预测结果的影响 :-)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
